﻿#include "common.h"
#include "MusicReader.h"

MyFilter my_filters;
InPCMFormat in_pcm_format(2, 44100, "s16");
bool is_pcm_format_set_successfully = false;
int open_file_error = 0;

InPCMFormat::InPCMFormat(int channels,
    int sample_rate,
    std::string sample_fmt)
    : pcm_sample_rate(sample_rate)
    , pcm_sample_fmt(AV_SAMPLE_FMT_S16)
    , pcm_layout(AV_CHANNEL_LAYOUT_STEREO)
{
    switch (channels) {
    case 1:
        pcm_layout = AV_CHANNEL_LAYOUT_MONO;
        break;
    case 2:
        pcm_layout = AV_CHANNEL_LAYOUT_STEREO;
        break;
    default:
        qDebug() << "Unsupported number of channels";
        throw std::runtime_error("Unsupported number of channels");
    }

    if (!sample_fmt.compare("u8"))
        pcm_sample_fmt = AV_SAMPLE_FMT_U8;
    else if (!sample_fmt.compare("s16"))
        pcm_sample_fmt = AV_SAMPLE_FMT_S16;
    else if (!sample_fmt.compare("s32"))
        pcm_sample_fmt = AV_SAMPLE_FMT_S32;
    else if (!sample_fmt.compare("f32"))
        pcm_sample_fmt = AV_SAMPLE_FMT_FLT;
}

void InPCMFormat::set_in_pcm_format(
    int channels,
    int sample_rate,
    std::string sample_fmt)
{
    channels == 1 ? pcm_layout = AV_CHANNEL_LAYOUT_MONO : pcm_layout = AV_CHANNEL_LAYOUT_STEREO;
    pcm_sample_rate = sample_rate;
    if (sample_fmt == "u8")
        pcm_sample_fmt = AV_SAMPLE_FMT_U8;
    else if (sample_fmt == "s16")
        pcm_sample_fmt = AV_SAMPLE_FMT_S16;
    else if (sample_fmt == "s32")
        pcm_sample_fmt = AV_SAMPLE_FMT_S32;
    else if (sample_fmt == "f32")
        pcm_sample_fmt = AV_SAMPLE_FMT_FLT;
    else {
        // 如果遇到未预期的采样格式，则抛出异常并设置为默认采样格式
        throw std::invalid_argument("Unsupported sample format: " + sample_fmt);
        pcm_sample_fmt = AV_SAMPLE_FMT_NONE; // 设置为默认不支持的采样格式
    }
}

std::string InPCMFormat::format_string()
{
    return std::string("channels: " + std::to_string(pcm_layout.nb_channels) + "    sample_rate: " + std::to_string(pcm_sample_rate) + "    sample_fmt: " + av_get_sample_fmt_name(pcm_sample_fmt));
}


MusicReader::MusicReader(const std::string& filename)
    : p_format_context(nullptr)
    //, per_channel_out_samples_count(-1)
    , p_codec_context(nullptr)
    , p_audio_stream(nullptr)
    , p_swr_context_SoundTouch(nullptr)
    , my_stereo_layout(AV_CHANNEL_LAYOUT_MONO)
    , my_out_sample_rate(0)
    , my_out_sample_fmt(AV_SAMPLE_FMT_NONE)
    , audio_stream_index(-1)
    , filter_graph(nullptr)
    , buffersrc_ctx(nullptr)
    , buffersink_ctx(nullptr)
    , is_pcm(false)
    //, current_frame_position(0)
{
    /*if (is_pcm)
        ;*/

    if (init_ffmpeg_audio(filename.c_str()) < 0) {
        free_ffmpeg();
        return;
        //throw std::runtime_error("FFmpeg initialization failed");
    }

    init_filters();

    // 如果需要，可以在这里初始化滤波器参数
    // 如果初始化失败，则在这里重新抛出异常
    // 请确保在这里正确初始化滤波器
    // this->adjust_filters(5.0f, 0.0f, -5.0f); // 调整滤波器，例如音频的高低频
}

MusicReader::~MusicReader()
{
    //avformat_network_deinit();
    free_ffmpeg();
}

void MusicReader::free_ffmpeg()
{
    if (p_format_context) {
        avformat_close_input(&p_format_context);
        p_format_context = nullptr;
    }
    if (p_codec_context) {
        avcodec_free_context(&p_codec_context);
        p_codec_context = nullptr;
    }
    if (p_swr_context_SoundTouch) {
        swr_free(&p_swr_context_SoundTouch);
        p_swr_context_SoundTouch = nullptr;
    }
    if (p_swr_context_output) {
        swr_free(&p_swr_context_output);
        p_swr_context_output = nullptr;
    }
    if (p_swr_context_SoundTouch) {
        swr_free(&p_swr_context_SoundTouch);
        p_swr_context_SoundTouch = nullptr;
    }

    // 释放 filter graph
    free_filters();

    /*if (p_audio_stream) delete p_audio_stream;*/
    avformat_network_deinit();
}

void MusicReader::set_pcm_format(AVCodecContext* p_codec_context)
{
    // 设置 AVFormatContext 格式
    p_format_context->iformat = av_find_input_format("s16le");
    if (!p_format_context->iformat)
        throw std::runtime_error("Input format not found");
}

//void MusicReader::set_pcm_format(AVCodecContext* p_codec_context) {
//    p_codec_context->request_sample_fmt = in_pcm_format.pcm_sample_fmt;
//    p_codec_context->ch_layout = in_pcm_format.pcm_layout;
//    p_codec_context->sample_rate = in_pcm_format.pcm_sample_rate;
//}

char MusicReader::init_ffmpeg_audio(const std::string& filename)
{
    // Initialize ffmpeg
    // This must be done before opening the file, otherwise it will fail
    free_ffmpeg();

    qDebug() << "Opening file: " << filename.c_str();

    // Initialize ffmpeg, older versions of ffmpeg may require additional initialization
    //av_register_all();

    // Initialize network, if needed
    avformat_network_init();

    // initialize AVCodecContext
    p_codec_context = avcodec_alloc_context3(nullptr);
    if (!p_codec_context)
        throw std::runtime_error("Could not allocate codec context");

    //if (filename.end())
    p_codec_context->sample_fmt = AV_SAMPLE_FMT_S16;
    p_codec_context->ch_layout.nb_channels = AV_CH_LAYOUT_STEREO;
    p_codec_context->sample_rate = 44100; // 设置默认值

    /*uint num = ~0;
    qDebug() << "num: " << num;*/

    p_format_context = avformat_alloc_context();

    // Open file
    // 设置音频格式为 pcm 并配置 AVFormatContext 的参数
    if (filename.rfind(".pcm") != std::string::npos)
    {
        is_pcm = true;

        if (in_pcm_format.pcm_layout.nb_channels == 1 &&
            in_pcm_format.pcm_sample_fmt == AV_SAMPLE_FMT_U8) {
            p_format_context->iformat = av_find_input_format("u8");
            //p_format_context->iformat->flags |= AVFMT_NOFILE;
        }
        else if ((in_pcm_format.pcm_layout.nb_channels == 1 &&
            in_pcm_format.pcm_sample_fmt == AV_SAMPLE_FMT_S16) ||
            (in_pcm_format.pcm_layout.nb_channels == 2
                && in_pcm_format.pcm_sample_fmt == AV_SAMPLE_FMT_U8))
            p_format_context->iformat = av_find_input_format("s16le");
        else if (in_pcm_format.pcm_layout.nb_channels == 2 &&
            in_pcm_format.pcm_sample_fmt == AV_SAMPLE_FMT_S16)
            p_format_context->iformat = av_find_input_format("s32le");
        else if (in_pcm_format.pcm_layout.nb_channels == 1 &&
            in_pcm_format.pcm_sample_fmt == AV_SAMPLE_FMT_S32)
            p_format_context->iformat = av_find_input_format("s32le");
        else if (in_pcm_format.pcm_layout.nb_channels == 2 &&
            in_pcm_format.pcm_sample_fmt == AV_SAMPLE_FMT_S32)
            p_format_context->iformat = av_find_input_format("s64le");
        else if (in_pcm_format.pcm_layout.nb_channels == 1 &&
            in_pcm_format.pcm_sample_fmt == AV_SAMPLE_FMT_FLT)
            p_format_context->iformat = av_find_input_format("f32le");
        else if (in_pcm_format.pcm_layout.nb_channels == 2 &&
            in_pcm_format.pcm_sample_fmt == AV_SAMPLE_FMT_FLT)
            p_format_context->iformat = av_find_input_format("f64le");
        else
            throw std::runtime_error("Unsupported PCM format");
        
        // 将音频格式设置为 pcm 并配置 AVFormatContext 的参数
        AVDictionary* options = nullptr;
        av_dict_set_int(&options, "sample_rate", in_pcm_format.pcm_sample_rate, 0);
        av_dict_set_int(&options, "channels", in_pcm_format.pcm_layout.nb_channels, 0);

        // 设置音频格式的字符串
        const char* fmt_str = nullptr;
        switch (in_pcm_format.pcm_sample_fmt) {
            case AV_SAMPLE_FMT_U8:
                fmt_str = "u8";
                break;
            case AV_SAMPLE_FMT_S16:
                fmt_str = "s16le";
                break;
            case AV_SAMPLE_FMT_S32:
                fmt_str = "s32le";
                break;
            case AV_SAMPLE_FMT_FLT:
                fmt_str = "f32le";
                break;
            default:
                fmt_str = nullptr;
                break;
        }
        if (fmt_str)
            av_dict_set(&options, "sample_fmt", fmt_str, 0);

        if ((open_file_error = avformat_open_input(&p_format_context, filename.c_str(), p_format_context->iformat, &options)) != 0) {
            //emit open_pcm_file_error();
            is_pcm_format_set_successfully = false;
            char errbuf[256];
            av_strerror(open_file_error, errbuf, sizeof(errbuf));
            qDebug() << "Could not open input file: " << errbuf;
            av_dict_free(&options);
            avcodec_free_context(&p_codec_context);
            avformat_network_deinit();
            return -1;
            //throw std::runtime_error("Could not open input file");
        }

        is_pcm_format_set_successfully = true;
        av_dict_free(&options);
    }
    else if ((open_file_error = avformat_open_input(&p_format_context, filename.c_str(), nullptr, nullptr)) != 0) {
        char errbuf[256];
        av_strerror(open_file_error, errbuf, sizeof(errbuf));
        qDebug() << "Could not open input file: " << errbuf;
        avcodec_free_context(&p_codec_context);
        avformat_network_deinit();
        return -1;
        //throw std::runtime_error("Could not open input file");
    }

    // Retrieve stream information --- 获取流信息
    if (avformat_find_stream_info(p_format_context, nullptr) < 0)
        throw std::runtime_error("Could not find stream information");

    // Find the first audio stream --- 查找第一个音频流
    audio_stream_index = av_find_best_stream(
        p_format_context,
        AVMEDIA_TYPE_AUDIO,
        -1, -1,
        &p_codec_context->codec,
        0);
    if (audio_stream_index < 0)
        throw std::runtime_error("Could not find audio stream");
    // 获取音频流
    p_audio_stream = p_format_context->streams[audio_stream_index];
    // 初始化 AVCodecContext
    p_codec_context = avcodec_alloc_context3(nullptr);
    if (!p_codec_context)
        throw std::runtime_error("无法配置 AVCodecContext");
    // 将音频流的参数复制到 AVCodecContext
    if (avcodec_parameters_to_context(p_codec_context, p_audio_stream->codecpar) < 0)
        throw std::runtime_error("无法将音频流的参数复制到 AVCodecContext");
    // 查找解码器
    const AVCodec* p_codec = avcodec_find_decoder(p_codec_context->codec_id);
    if (!p_codec)
        throw std::runtime_error("无法找到解码器实现");
    // 打开解码器
    if (avcodec_open2(p_codec_context, p_codec, nullptr) < 0)
        throw std::runtime_error("无法打开解码器");

    //av_channel_layout_default(&my_stereo_layout, AV_CH_LAYOUT_STEREO); // 查找默认声道

    reset_audio_sampler_SoundTouch();
    reset_audio_sampler_output();

    /*p_codec_context->bit_rate = 128000;
    qDebug() << "Bitrate set to: " << p_codec_context->bit_rate;*/

    audio_filename = filename;
	return 0;
}

void MusicReader::reset_audio_sampler_SoundTouch(
    AVChannelLayout st_layout,
    int sample_rate,
    AVSampleFormat sample_fmt)
{
    std::lock_guard<std::mutex> lock(reader_mutex);
    if (p_swr_context_SoundTouch)
        swr_free(&p_swr_context_SoundTouch);

    /*if (is_pcm) {
        stereo_layout = in_pcm_format.pcm_layout;
        out_sample_rate = in_pcm_format.pcm_sample_rate;
        out_sample_fmt = in_pcm_format.pcm_sample_fmt;
    }*/

    p_swr_context_SoundTouch = swr_alloc();
    if (!p_swr_context_SoundTouch)
        throw std::runtime_error("Could not allocate resampler context");

    // 输入
    av_opt_set_chlayout(p_swr_context_SoundTouch, "in_chlayout", &p_codec_context->ch_layout, 0);
    av_opt_set_int(p_swr_context_SoundTouch, "in_sample_rate", p_codec_context->sample_rate, 0);
    av_opt_set_sample_fmt(p_swr_context_SoundTouch, "in_sample_fmt", p_codec_context->sample_fmt, 0);
    
    // 输出
    st_stereo_layout = st_layout;
    av_opt_set_chlayout(p_swr_context_SoundTouch, "out_chlayout", &st_layout, 0);
    st_out_sample_rate = sample_rate;
    av_opt_set_int(p_swr_context_SoundTouch, "out_sample_rate", sample_rate, 0);
    st_out_sample_fmt = sample_fmt;
    av_opt_set_sample_fmt(p_swr_context_SoundTouch, "out_sample_fmt", sample_fmt, 0);

    int ret = swr_init(p_swr_context_SoundTouch);
    if (ret < 0) {
        char errbuf[128];
        av_strerror(ret, errbuf, sizeof(errbuf));
        qDebug() << "Could not initialize sampler: " << errbuf;
        throw std::runtime_error("Could not initialize sampler");
    }
}

void MusicReader::reset_audio_sampler_output(
    AVChannelLayout out_layout,
    int out_sample_rate,
    AVSampleFormat out_sample_fmt)
{
    std::lock_guard<std::mutex> lock(reader_mutex);
    if (p_swr_context_output)
        swr_free(&p_swr_context_output);

    /*if (is_pcm) {
        stereo_layout = in_pcm_format.pcm_layout;
        out_sample_rate = in_pcm_format.pcm_sample_rate;
        out_sample_fmt = in_pcm_format.pcm_sample_fmt;
    }*/

    p_swr_context_output = swr_alloc();
    if (!p_swr_context_output)
        throw std::runtime_error("Could not allocate resampler context");

    // 输入
    av_opt_set_chlayout(p_swr_context_output, "in_chlayout", &st_stereo_layout, 0);
    av_opt_set_int(p_swr_context_output, "in_sample_rate", st_out_sample_rate, 0);
    av_opt_set_sample_fmt(p_swr_context_output, "in_sample_fmt", AV_SAMPLE_FMT_FLTP, 0); // 滤镜使用后变成了 AV_SAMPLE_FMT_FLTP 格式

    // 输出
    av_opt_set_chlayout(p_swr_context_output, "out_chlayout", &out_layout, 0);
    my_stereo_layout = out_layout;
    av_opt_set_int(p_swr_context_output, "out_sample_rate", out_sample_rate, 0);
    my_out_sample_rate = out_sample_rate;
    av_opt_set_sample_fmt(p_swr_context_output, "out_sample_fmt", out_sample_fmt, 0);
    my_out_sample_fmt = out_sample_fmt;

    int ret = swr_init(p_swr_context_output);
    if (ret < 0) {
        char errbuf[128];
        av_strerror(ret, errbuf, sizeof(errbuf));
        qDebug() << "Could not initialize sampler: " << errbuf;
        throw std::runtime_error("Could not initialize sampler");
    }
}

const double MusicReader::get_duration()
{
    if (!p_format_context || !p_audio_stream)
        throw std::runtime_error("File not opened");
    // The duration is stored in microseconds, so maybe we need convert it to seconds
    double duration_seconds;
    /*if (p_audio_stream->duration != AV_NOPTS_VALUE) {
        duration_seconds = p_audio_stream->duration * av_q2d(p_audio_stream->time_base);
        qDebug() << "Audio duration: " << duration_seconds << " seconds";
    }
    else {
        qDebug() << "Audio duration not available";*/

        // 计算总时长
        int64_t total_duration = 0;
        AVPacket pkt;
        while (av_read_frame(p_format_context, &pkt) >= 0) {
            if (pkt.stream_index == p_audio_stream->index) {
                total_duration += pkt.duration;
            }
            av_packet_unref(&pkt);
        }
        duration_seconds = total_duration * av_q2d(p_audio_stream->time_base);
        qDebug() << "Calculated audio duration: " << duration_seconds << " seconds";
    //}

    seek(0); // 可移动到上次停止的位置
    //return p_format_context->duration; // / static_cast<double>(AV_TIME_BASE); 音频中储存的
    return duration_seconds; // 手动计算的
}

void MusicReader::save_resampled_audio(const char* output_filename)
{
    // 保存为 16 位 PCM 格式
    // TODO: Implement this function
    if (!p_format_context || !p_audio_stream)
        throw std::runtime_error("File not opened");

    // Open output file in binary write mode
    FILE* output_file;
    fopen_s(&output_file, output_filename, "wb");
    if (!output_file)
        throw std::runtime_error("Could not open output file");

    // Process each frame
    AVFrame* p_save_frame;
    //p_save_frame = av_frame_alloc();
    seek(0);
    while ((p_save_frame = next_audio_out_frame()) != nullptr)
    {
        qDebug() << "Frame: " << p_save_frame;
        if (p_save_frame->format != AV_SAMPLE_FMT_S16 ||
            p_save_frame->sample_rate != 44100 ||
            p_save_frame->ch_layout.nb_channels != 2)
        {
            av_frame_unref(p_save_frame);
            av_frame_free(&p_save_frame);
            throw std::runtime_error("Frame format does not match expected format");
        }

        // Write raw PCM data to file
        size_t data_size = av_get_bytes_per_sample((AVSampleFormat)p_save_frame->format) *
            p_save_frame->nb_samples * p_save_frame->ch_layout.nb_channels;
        if (fwrite(p_save_frame->data[0], 1, data_size, output_file) != data_size)
        {
            av_frame_unref(p_save_frame);
            av_frame_free(&p_save_frame);
            fclose(output_file);
            throw std::runtime_error("Error writing raw PCM data to file");
        }

        av_frame_unref(p_save_frame);
        av_frame_free(&p_save_frame);
    }

    // Close output file
    fclose(output_file);
}

void MusicReader::seek(int64_t pos)
{
    if (!p_format_context || !p_codec_context || !p_audio_stream)
        throw std::runtime_error("File not opened");

    std::lock_guard<std::mutex> lock(reader_mutex);
    // 将目标时间位置转换为以 time_base 为单位的值
    int64_t timestamp = av_rescale_q(pos, AV_TIME_BASE_Q, p_audio_stream->time_base);
    if (av_seek_frame(p_format_context, audio_stream_index, timestamp, AVSEEK_FLAG_BACKWARD) < 0)
        throw std::runtime_error("failed to seek");

    // 刷新解码器缓冲区
    avcodec_flush_buffers(p_codec_context);

    /*qDebug() << "Seek to position: " << pos << " (" << timestamp << " microseconds)";*/

    // 丢弃 seek 之前的帧，只保留定位后的新帧
    //AVPacket packet;
    /*av_init_packet(&packet);
    packet.data = nullptr;
    packet.size = 0;*/

    //bool found = false;
    //while (av_read_frame(p_format_context, &packet) >= 0) {
    //    if (packet.stream_index == audio_stream_index) {
    //        int ret = avcodec_send_packet(p_codec_context, &packet);
    //        if (ret < 0) {
    //            av_packet_unref(&packet);
    //            continue;
    //        }
    //        AVFrame* temp_frame = av_frame_alloc();
    //        ret = avcodec_receive_frame(p_codec_context, temp_frame);
    //        if (ret >= 0) {
    //            // 找到 seek 后的第一帧
    //            current_frame_position = temp_frame->best_effort_timestamp;
    //            av_frame_free(&temp_frame);
    //            found = true;
    //            av_packet_unref(&packet);
    //            break;
    //        }
    //        av_frame_free(&temp_frame);
    //    }
    //    av_packet_unref(&packet);
    //}
    //av_packet_unref(&packet);
    //if (!found) {
    //    throw std::runtime_error("Seek 到指定位置失败");
    //}
}

AVFrame* MusicReader::next_audio_out_frame()
{
    // 返回下一帧的音频，记得释放输出音频帧内存
    std::lock_guard<std::mutex> lock(reader_mutex);
    static int64_t audio_pts = 0;
    static AVPacket* cached_packet = nullptr; // 静态变量保存未完全送入的 packet
    static bool packet_pending = false; // 是否有未完全送入的 packet

    AVFrame* p_frame = av_frame_alloc();
    if (!p_frame)
        throw std::runtime_error("Could not allocate audio frame");
    AVFrame* p_out_frame = av_frame_alloc();
    if (!p_out_frame)
        throw std::runtime_error("Could not allocate output frame");
    p_out_frame->format = st_out_sample_fmt; // float for SoundTouch
    p_out_frame->ch_layout = st_stereo_layout;
    p_out_frame->sample_rate = st_out_sample_rate;

    int ret = 0;
    while (true) {
        // 1. 先尝试 receive frame
        ret = avcodec_receive_frame(p_codec_context, p_frame);
        if (ret == 0) {
            // 解码成功，重采样
            int delay = swr_get_delay(p_swr_context_SoundTouch, p_frame->sample_rate);
            int out_sample_count = av_rescale_rnd(
                delay + p_frame->nb_samples,
                st_out_sample_rate,
                p_frame->sample_rate,
                AV_ROUND_UP);
            p_out_frame->nb_samples = out_sample_count;
            if (av_frame_get_buffer(p_out_frame, 0) < 0)
                throw std::runtime_error("Could not allocate output frame data buffer");
            // 重采样到 float 格式，供 SoundTouch 处理
            int conv_ret = swr_convert(p_swr_context_SoundTouch,
                p_out_frame->data,
                out_sample_count,
                p_frame->data,
                p_frame->nb_samples);
            if (conv_ret < 0) {
                break;
            }
            p_out_frame->nb_samples = conv_ret;

            if (p_frame->pts != AV_NOPTS_VALUE && p_frame->pts >= 0) {
                p_out_frame->pts = p_frame->pts;
                audio_pts = p_frame->pts + p_out_frame->nb_samples;
            }
            else {
                // 估算 pts
                p_out_frame->pts = audio_pts;
                audio_pts += p_out_frame->nb_samples;
            }

            av_frame_unref(p_frame);
            av_frame_free(&p_frame);
            // 注意：不要释放p_out_frame，返回给调用者
            return p_out_frame;
        }
        else if (ret == AVERROR(EAGAIN)) { /*需要 send packet*/ }
        else if (ret == AVERROR_EOF || ret < 0) {
            break;
        }

        // 2. send packet
        if (!packet_pending) {
            // 没有未送入的 packet，读取新 packet
            AVPacket* new_packet = av_packet_alloc();
            while (av_read_frame(p_format_context, new_packet) >= 0) {
                if (new_packet->stream_index != audio_stream_index) {
                    av_packet_unref(new_packet);
                    continue;
                }
                // 找到音频流packet
                cached_packet = new_packet;
                packet_pending = true;
                break;
            }
            if (!packet_pending) {
                // 没有更多packet可读
                av_packet_free(&new_packet);
                break;
            }
        }
        // 尝试送入packet
        ret = avcodec_send_packet(p_codec_context, cached_packet);
        if (ret == AVERROR(EAGAIN)) {
            // 解码器满了，先 receive frame，下次再 send
            // 不释放 cached_packet
            continue;
        }
        else {
            // send成功或其他错误，都释放 cached_packet
            av_packet_unref(cached_packet);
            av_packet_free(&cached_packet);
            packet_pending = false;
            if (ret < 0) {
                // 其他错误，尝试下一个 packet
                continue;
            }
        }
    }
    // 释放资源
    if (packet_pending && cached_packet) {
        av_packet_unref(cached_packet);
        av_packet_free(&cached_packet);
        packet_pending = false;
    }
    av_frame_free(&p_frame);
    av_frame_free(&p_out_frame);
    return nullptr;
}
